# 数据加载与预处理
import pandas as pd

def sample_data(data1):
    # 计算用户活跃度
    user_activity = data1['user'].value_counts()
    active_users = user_activity[user_activity >= 20].index
    inactive_users = user_activity[user_activity < 20].index
    
    # 计算图书活跃度
    book_activity = data1['item'].value_counts()
    active_books = book_activity[book_activity >= 10].index
    inactive_books = book_activity[book_activity < 10].index

    # 计算每个类别的比例
    active_user_ratio = len(active_users) / len(user_activity)
    inactive_user_ratio = len(inactive_users) / len(user_activity)
    active_book_ratio = len(active_books) / len(book_activity)
    inactive_book_ratio = len(inactive_books) / len(book_activity)
    
    # 根据总数来确定抽样的用户数和图书数
    total_users = 2000
    total_books = 100000
    
    # 分层抽样：根据比例抽取样本
    active_user_sample_size = min(len(active_users), int(total_users * active_user_ratio))
    inactive_user_sample_size = min(len(inactive_users), int(total_users * inactive_user_ratio))
    
    active_book_sample_size = min(len(active_books), int(total_books * active_book_ratio))
    inactive_book_sample_size = min(len(inactive_books), int(total_books * inactive_book_ratio))
    
    active_user_sample = data1[data1['user'].isin(active_users)].sample(n=active_user_sample_size, random_state=42) if active_user_sample_size > 0 else data1[data1['user'].isin(active_users)]
    inactive_user_sample = data1[data1['user'].isin(inactive_users)].sample(n=inactive_user_sample_size, random_state=42) if inactive_user_sample_size > 0 else data1[data1['user'].isin(inactive_users)]
    
    active_book_sample = data1[data1['item'].isin(active_books)].sample(n=active_book_sample_size, random_state=42) if active_book_sample_size > 0 else data1[data1['item'].isin(active_books)]
    inactive_book_sample = data1[data1['item'].isin(inactive_books)].sample(n=inactive_book_sample_size, random_state=42) if inactive_book_sample_size > 0 else data1[data1['item'].isin(inactive_books)]
    # 合并数据
    final_data = pd.concat([active_user_sample, inactive_user_sample, active_book_sample, inactive_book_sample])

    # 通过交集确保只保留用户和图书的交互数据
    sampled_users = final_data['user'].drop_duplicates()
    sampled_books = final_data['item'].drop_duplicates()

    final_data = final_data[final_data['user'].isin(sampled_users) & final_data['item'].isin(sampled_books)]

    # 确保数据集不超过1000个用户和1000个物品
    # final_data = final_data.sample(n=1000, random_state=42)

    print(f"Final sampled data shape: {final_data.shape}")
    
    print(final_data)
    
    return final_data

def load_data():
    print("start loading...")
    f1 = "amazon_book\Books.csv"
    f2 = "book_crossing\Ratings_modified.csv"
    f3 = "goodreads\goodreads_interactions.csv"
    print("loading amazon_book dataset...")
    data1 = pd.read_csv(f1, on_bad_lines='skip')
    print("loading book_crossing dataset...")
    data2 = pd.read_csv(f2, on_bad_lines='skip')
    print("loading goodreads dataset...")
    data3 = pd.read_csv(f3, on_bad_lines='skip')
    # 删除字段数超过3个或不等于3个的行
    data1 = data1[data1.notnull().sum(axis=1) == 3]  # 保留字段数为3的行
    data2 = data2[data2.notnull().sum(axis=1) == 3]  # 保留字段数为3的行
    data3 = data3[data3.notnull().sum(axis=1) == 3]  # 保留字段数为3的行
    
    # 将book_crossing的评分从0-10映射到0-5
    data2['rating'] = (data2['rating'] / 10) * 5
    # 使用格式化字符串打印数据集的形状
    print(f"data1 shape: {data1.shape}, data2 shape: {data2.shape}, data3 shape: {data3.shape}")
    # 查看每个数据集的用户数和图书数
    users1 = len(data1['user'].unique())
    items1 = len(data1['item'].unique())
    
    users2 = len(data2['user'].unique())
    items2 = len(data2['item'].unique())
    
    users3 = len(data3['user'].unique())
    items3 = len(data3['item'].unique())
    
    print(f"amazon_book dataset - Users: {users1}, Books: {items1}")
    print(f"book_crossing dataset - Users: {users2}, Books: {items2}")
    print(f"goodreads dataset - Users: {users3}, Books: {items3}")
    # 将数据映射成字典
    data = {
        "a": sample_data(data1),
        "b": sample_data(data2),
        "g": sample_data(data3)
    }
    print("complete loading...")
    return data


